# 时域有限差分求解三维电磁场
# 使用周期性边界条件
# Gitee Repo

import numpy as np
import torch
import scipy

torch.set_default_device('cuda' if torch.cuda.is_available() else 'cpu')

L = 2
dx = 0.05
dt = 0.001
c = 1

sigma = (c*dt/dx)**2

x,y,z = torch.meshgrid(torch.arange(-L,L,dx),torch.arange(-L,L,dx),torch.arange(-L,L,dx))
n = x.shape[0]

E0 = torch.zeros((n,n,n,4))
B0 = torch.zeros((n,n,n,4))
E1 = torch.zeros((n,n,n,4))
B1 = torch.zeros((n,n,n,4))

TICK = 3000
for tick in range(TICK+2):
    # \curl: yz-zy,zx-xz,xy-yx
    # \pdv{E}{t} = 1/(\mu_0 \epsilon_0) \curl B - 1/(\epsilon_0) j
    E1[:,:,:,1] = E0[:,:,:,1] + dt/dx * ((B0[:,:,:,3] - torch.roll(B0[:,:,:,3],-1,dims=2-1)) - (B0[:,:,:,2] - torch.roll(B0[:,:,:,2],-1,dims=3-1)))
    E1[:,:,:,2] = E0[:,:,:,2] + dt/dx * ((B0[:,:,:,1] - torch.roll(B0[:,:,:,1],-1,dims=3-1)) - (B0[:,:,:,3] - torch.roll(B0[:,:,:,3],-1,dims=1-1)))
    E1[:,:,:,3] = E0[:,:,:,3] + dt/dx * ((B0[:,:,:,2] - torch.roll(B0[:,:,:,2],-1,dims=1-1)) - (B0[:,:,:,1] - torch.roll(B0[:,:,:,1],-1,dims=2-1)))

    E1[int(n/2),int(n/2),int(n/2),3] += dt/dx**3*np.sin(2*np.pi/1*tick*dt)
    
    # \pdv{B}{t} = - \curl E
    B1[:,:,:,1] = B0[:,:,:,1] - dt/dx * ((torch.roll(E1[:,:,:,3],+1,dims=2-1) - E1[:,:,:,3]) - (torch.roll(E1[:,:,:,2],+1,dims=3-1) - E1[:,:,:,2]))
    B1[:,:,:,2] = B0[:,:,:,2] - dt/dx * ((torch.roll(E1[:,:,:,1],+1,dims=3-1) - E1[:,:,:,1]) - (torch.roll(E1[:,:,:,3],+1,dims=1-1) - E1[:,:,:,3]))
    B1[:,:,:,3] = B0[:,:,:,3] - dt/dx * ((torch.roll(E1[:,:,:,2],+1,dims=1-1) - E1[:,:,:,2]) - (torch.roll(E1[:,:,:,1],+1,dims=2-1) - E1[:,:,:,1]))

    print(tick)
    if tick >= TICK:
        break

    E0,E1=E1,E0
    B0,B1=B1,B0

scipy.io.savemat('EB.mat',{'L':L, 'x':x.cpu().numpy(),'y':y.cpu().numpy(),'z':z.cpu().numpy(),'E':E1.cpu().numpy(), 'B':B1.cpu().numpy()},do_compression=True)

